#define _GNU_SOURCE 1
#include <sched.h> //for binding process to processor
#include <stdio.h>
#include <stdlib.h>
#include <pthread.h>
#include <errno.h>
#include <time.h>
#include <sys/wait.h>
#include <sys/time.h>
#include <unistd.h>
#include "blocking_queue.h"
#include <unistd.h>

#define WORK_DELAY 1

pthread_barrier_t barr;
queue_t queue;

unsigned long* producers_results;
unsigned long* consumers_results;

int *producers_cpus;
int *consumers_cpus;
int num_of_cpus = 1;

int volatile stop = 0;

void handle_affinity(char* cpu_arg){
		int pid = getpid();
		if (cpu_arg != NULL) {
		fprintf(stderr, "cpu arg %s\n", cpu_arg);
		cpu_set_t* cpu_set;
		char* p = NULL;
		int num_of_cpus = 0;
		if ((p = strtok(cpu_arg, ":")) == NULL) {
			perror("split error");
			return;
		} else {
			num_of_cpus = atoi(p);
			fprintf(stderr, "num of cpus %d\n", num_of_cpus);
			cpu_set = CPU_ALLOC(num_of_cpus);
		}
		size_t size = CPU_ALLOC_SIZE(num_of_cpus);

		CPU_ZERO_S(size, cpu_set);
		while ((p = strtok(NULL, ",")) != NULL) {
			int cpu_num = atoi(p);
			fprintf(stderr, "cpu_num to run on %d\n", cpu_num);
			CPU_SET_S(cpu_num, size, cpu_set);
		}

		fflush(stderr);

		if (sched_setaffinity(pid, size, cpu_set) < 0) {
			perror("sched_setaffinity");
		}

		if (sched_getaffinity(pid, size, cpu_set) < 0) {
			perror("sched_getaffinity");
			return -1;
		}
		fprintf(stderr, "pid %d's affinity mask: %d\n", pid, *cpu_set);
		CPU_FREE(cpu_set);
	}
}

void* produce(void* arg) {
	int pnum = *((int*) arg);
	free((int*) arg);
	pthread_barrier_wait(&barr);

	unsigned long i = 0;
	while (!stop) {
		int* data = malloc(sizeof(int));
		data[0] = i;
		if (queue_put(&queue, data) == 0) {
			i++;
		}
	}
	producers_results[pnum] = i;
	pthread_exit(0);
}

void* consume(void* arg) {
	int cnum = *((int*) arg);
	free((int*) arg);
	pthread_barrier_wait(&barr);

	unsigned long i = 0;
	while (!stop) {
		int* data = queue_get(&queue);
		//dummy work
		sleep(WORK_DELAY);
		if (data != NULL) {
			i++;
		}
		free(data);
	}
	consumers_results[cnum] = i;
	pthread_exit(0);
}

void join(pthread_t threads[], int size) {
	void* status;
	int t, rc;
	for (t = 0; t < size; t++) {
		rc = pthread_join(threads[t], &status);
		if (rc) {
			printf("ERROR; return code from pthread_join() is %d\n", rc);
			exit(-1);
		}
	}
	fflush(stdout);
}

int main(int argc, char *argv[]) {
	if (argc != 7) {
		fprintf(
				stderr,
				"usage: producer-consumer <#producers> <#consumers> <time to run> <queue length> <timeout> <cpuNum:cpusToRunOn[,]>");
		exit(-1);
	}
	int num_producers = atoi(argv[1]);
	int num_consumers = atoi(argv[2]);
	int time_to_run = atoi(argv[3]);
	long queue_length = atol(argv[4]);
	int queue_timeout = atoi(argv[5]);
	handle_affinity(argv[6]);

	pthread_t* producers = malloc(sizeof(pthread_t) * num_producers);
	producers_results = malloc(sizeof(unsigned long) * num_producers);
	pthread_t* consumers = malloc(sizeof(pthread_t) * num_consumers);
	consumers_results = malloc(sizeof(unsigned long) * num_consumers);

	struct timeval startTime;
	struct timeval endTime;

	pthread_attr_t attr;
	if (pthread_attr_init(&attr)) {
		printf("Could not create a attr\n");
		exit(-1);
	}
	if (pthread_attr_setdetachstate(&attr, PTHREAD_CREATE_JOINABLE)) {
		printf("Could not set attr\n");
		exit(-1);
	}

	if (pthread_barrier_init(&barr, NULL, num_producers + num_consumers)) {
		printf("Could not create a barrier\n");
		exit(-1);
	}
	queue_init(&queue, queue_length, queue_timeout);

	int rc;
	long t;
	for (t = 0; t < num_producers; t++) {
		int* arg = malloc(sizeof(int));
		*arg = t;
		rc = pthread_create(&producers[t], &attr, produce, arg);
		if (rc) {
			printf("ERROR; return code from pthread_create() is %d\n", rc);
			exit(-1);
		}
	}
	for (t = 0; t < num_consumers; t++) {
		int* arg = malloc(sizeof(int));
		*arg = t;
		rc = pthread_create(&consumers[t], &attr, consume, arg);
		if (rc) {
			printf("ERROR; return code from pthread_create() is %d\n", rc);
			exit(-1);
		}
	}

	gettimeofday(&startTime, NULL);
	int delay = 0;
	do {
		gettimeofday(&endTime, NULL);
		delay = (endTime.tv_sec - startTime.tv_sec);
	} while (delay < time_to_run);
	stop = 1;

	join(producers, num_producers);
	join(consumers, num_consumers);

	unsigned long long psum = 0;
	for (t = 0; t < num_producers; t++) {
		psum += producers_results[t];
	}
	unsigned long long csum = 0;
	for (t = 0; t < num_consumers; t++) {
		csum += consumers_results[t];
	}

	printf("%llu:%llu\n", psum, csum);

	pthread_attr_destroy(&attr);
	pthread_barrier_destroy(&barr);
	queue_destroy(&queue);
	free(producers);
	free(consumers);
	pthread_exit(NULL);
}
